# -*- coding:utf-8 -*-
# author: lkz
# date: 2016/11/21 16:32

import socket
import struct


class ICMPSocket(object):
    def __init__(self, icmp_type, code, data_size):
        self._type = icmp_type
        self._code = code
        self._socket = socket.socket(socket.AF_INET, socket.SOCK_RAW, socket.IPPROTO_ICMP)
        self._socket.setblocking(0)

        pkt_size = data_size + 4
        self._pkt_packer = struct.Struct("!BBH%ss" % data_size)
        fmt = "!%sH%sB" % ((pkt_size / 2), (pkt_size & 1))
        self._double_bytes_parser = struct.Struct(fmt)

    def fileno(self):
        return self._socket.fileno()

    def close(self):
        self._socket.close()

    def __del__(self):
        self.close()

    def cal_checksum(self, pkt):
        pkt = bytes(pkt)
        double_bytes_list = self._double_bytes_parser.unpack(pkt)

        checksum = sum(double_bytes_list)
        checksum = (checksum >> 16) + (checksum & 0xffff)
        checksum += checksum >> 16

        return ~checksum & 0xffff

    def sendto(self, dst_ip, data):
        assert isinstance(data, str), 'data to be sent must be string type'

        _pkt = self._pkt_packer.pack(self._type, self._code, 0, data)
        checksum = self.cal_checksum(_pkt)
        pkt = self._pkt_packer.pack(self._type, self._code, checksum, data)

        n_bytes = self._socket.sendto(pkt, (dst_ip, 1))

        return n_bytes
